package zhengqc.graph.algorithm.mst;


import zhengqc.graph.algorithm.ConnectedComponent;
import zhengqc.graph.data_structure.WeightedEdge;
import zhengqc.graph.data_structure.WeightedGraph;
import zhengqc.union_find.UnionFind;
import zhengqc.union_find.UnionFindUsingQuickUnionPathCompression;

import java.util.ArrayList;
import java.util.Collections;

/* 暂时仅支持无向有权图 */
public class Kruskal {
    private WeightedGraph graph;
    private ArrayList<WeightedEdge> mst;
    private int minimalCost;

    public Kruskal(WeightedGraph graph) {
        this.graph = graph;
        this.minimalCost = 0;
        this.mst = new ArrayList<>();
        ConnectedComponent wcc = new ConnectedComponent(graph);
        if (wcc.getCount() > 1) {
            return;
        }
        ArrayList<WeightedEdge> edges = new ArrayList<>();
        for (int v = 0; v < graph.getV(); ++v) {
            for (int w: graph.getAdjacency(v)) {
                if (v < w) { // 保证每条边仅加入一次, 且满足 v < w
                    edges.add(new WeightedEdge(v, w, graph.getWeighted(v, w)));
                }
            }
        }
        Collections.sort(edges);
        UnionFind unionFind = new UnionFindUsingQuickUnionPathCompression(graph.getV());
        for (WeightedEdge edge: edges) {
            int v = edge.getV();
            int w = edge.getW();
            // 若 v 和 w 不在同一个连通分量, 则需要加边
            if (! unionFind.isConnected(v, w)) {
                mst.add(edge);
                unionFind.unionElement(v, w);
                minimalCost += edge.getWeight();
            }
        }
    }

    public ArrayList<WeightedEdge> getMST() {
        return this.mst;
    }

    public int getMinimalCost() {
        return minimalCost;
    }
}
